import json
import pandas as pd
import numpy as np
from langchain_openai import ChatOpenAI
from langchain.chains import LLMChain
from langchain_core.prompts import PromptTemplate
from sentence_transformers import SentenceTransformer
import torch


class HistoricalCaseRetriever:
    """历史用例检索器"""

    def __init__(self, case_db_path):
        self.case_db = self._load_case_database(case_db_path)
        self.model = SentenceTransformer('paraphrase-multilingual-MiniLM-L12-v2')
        self._create_semantic_index()

    def _load_case_database(self, path):
        """加载测试用例库"""
        # if path.endswith('.json'):
        #     with open(path, 'r', encoding='utf-8') as f:
        #         return json.load(f)
        # else:
                # 默认示例数据
        return [
                    {"title": "用户登录功能测试", "cases": [
                        "功能: 用户登录\n场景: 正确手机号登录\n当 输入有效的手机号和验证码\n那么 系统允许登录",
                        "功能: 用户登录\n场景: 验证码过期\n当 输入过期的验证码\n那么 系统提示验证码无效"
                    ]},
                    {"title": "订单支付功能测试", "cases": [
                        "功能: 订单支付\n场景: 微信支付成功\n当 选择微信支付并确认支付\n那么 订单状态变更为已支付",
                        "功能: 订单支付\n场景: 支付超时\n当 超过30分钟未支付\n那么 订单自动取消"
                    ]}
                ]

    def _create_semantic_index(self):
        """创建语义检索索引"""
        self.case_texts = []
        for category in self.case_db:
            for case in category['cases']:
                self.case_texts.append(f"{category['title']}: {case}")

        self.case_embeddings = self.model.encode(self.case_texts,
                                                 convert_to_tensor=True)

    def retrieve_similar_cases(self, requirement, top_k=3):
        """检索相似历史用例"""
        query_embedding = self.model.encode(requirement, convert_to_tensor=True)

        # 计算余弦相似度
        cos_scores = torch.nn.functional.cosine_similarity(
            query_embedding, self.case_embeddings)

        # 获取最相似的top_k个用例
        top_indices = torch.topk(cos_scores, k=top_k).indices.tolist()
        return [self.case_texts[idx] for idx in top_indices]


class TestCaseGenerator:
    """测试用例生成器"""

    def __init__(self, api_key, company_rules=None):
        self.llm = ChatOpenAI(model="gpt-3.5-turbo", api_key=api_key, temperature=0.7)
        self.rules = company_rules or self.default_rules()

        # 创建提示词模板
        self.prompt_template = PromptTemplate(
            input_variables=["requirement", "similar_cases", "rules"],
            template="""
            作为高级测试工程师，请根据需求生成符合规范的测试用例。

            # 公司测试规范
            {rules}

            # 需求描述
            {requirement}

            # 相关历史用例参考
            {similar_cases}

            # 生成要求
            1. 只输出测试用例，不需要需求描述
            2. 使用Gherkin语法：Feature, Scenario, Given, When, Then
            3. 格式：
               功能: [功能名称]
               场景: [场景描述]
               当[条件]
               且[条件]（可选）
               那么[结果]
            4. 生成3-5个核心场景
            """
        )

        # 创建处理链
        self.chain = LLMChain(llm=self.llm, prompt=self.prompt_template)

    def default_rules(self):
        """默认测试规范"""
        return {
            "priority": "必须包含优先级标记[P0-P3]",
            "coverage": "需包含正向、逆向和边界场景",
            "format": "使用标准Gherkin语法",
            "data": "参数化测试数据"
        }

    def generate(self, requirement, similar_cases):
        """生成测试用例"""
        # 格式化相似用例
        formatted_cases = "\n".join(similar_cases) if similar_cases else "无相似用例"

        # 生成用例
        response = self.chain.invoke({
            "requirement": requirement,
            "similar_cases": formatted_cases,
            "rules": json.dumps(self.rules, ensure_ascii=False)
        })

        return response['text']


class TestCaseFactory:
    """测试用例工厂（主入口）"""

    def __init__(self, api_key, case_db_path="test_cases.json", rules=None):
        self.retriever = HistoricalCaseRetriever(case_db_path)
        self.generator = TestCaseGenerator(api_key, rules)

    def generate_test_cases(self, requirement):
        """生成测试用例工作流"""
        # 检索相似用例
        similar_cases = self.retriever.retrieve_similar_cases(requirement)

        # 生成新用例
        test_cases = self.generator.generate(requirement, similar_cases)

        return test_cases


if __name__ == "__main__":
    # 配置信息
    API_KEY = "sk-079f9ad2ad3f457ebd6e6eb90f56fb53"  # 替换为你的OpenAI API密钥

    # 示例需求（使用发版值班规则）
    REQUIREMENT = """
    发版值班规则：
    1. 自动化失败后，值班测试需要拨打三次手机号码（不是企微）给测试owner
    2. 三次无响应后，必须向上级汇报（组测试负责人或BU技术负责人）
    3. 未经上级确认，不得擅自结束发布流程，否则需额外值班
    """

    # 创建工厂实例
    factory = TestCaseFactory(API_KEY)

    # 生成测试用例
    print("正在生成测试用例...")
    test_cases = factory.generate_test_cases(REQUIREMENT)

    print("\n生成的测试用例:")
    print(test_cases)